package LibDL.eval.cluster;

import LibDL.core.Tensor;
import LibDL.core.functional;

public class ARIEvaluator extends AbstractSupervisedClusterEvaluator{
    @Override
    protected double core(Tensor truth, Tensor pred) {
        checkClustering(truth, pred);

        int nSamples = truth.size(0);
        int nClasses = functional.unique(truth,true).get(0).size(0);
        int nCluster = functional.unique(pred,true).get(0).size(0);
        if ((nClasses == 1 && nCluster == 1)
                || (nClasses == 0 && nCluster == 0)
                || (nClasses == nSamples && nCluster == nSamples))
            return 1.;

        Tensor contingency = getContingencyMatrix();
        double sumCombC = contingency.sum(1).tolist_double().stream()
                .mapToDouble(d->comb2(d.intValue())).sum();;
        double sumCombK = contingency.sum(0).tolist_double().stream()
                .mapToDouble(d->comb2(d.intValue())).sum();
        double sumComb = contingency.tolist_double().stream()
                .mapToDouble(d->comb2(d.intValue())).sum();
        double probComb = sumCombC * sumCombK / comb2(nSamples);
        double meanComb = (sumCombC + sumCombK) / 2;

        return (sumComb - probComb) / (meanComb - probComb);
    }
}
